# Robust MEKF attitude estimation simulation with:
# - Magnetometer anomaly detection (angle-based + Mahalanobis gating)
# - Accelerometer linear-accel detection (norm deviation -> downweight)
# - Robust Mahalanobis outlier rejection for stacked measurement
# - Proper error reporting using rotation-vector (avoids Euler singularities)
#
# Run as-is. Requires: numpy, matplotlib

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(2)

# ---------- quaternion utilities ----------
def skew(v):
    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])

def quat_mul(q1, q2):
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    return np.array([
        w1*w2 - x1*x2 - y1*y2 - z1*z2,
        w1*x2 + x1*w2 + y1*z2 - z1*y2,
        w1*y2 - x1*z2 + y1*w2 + z1*x2,
        w1*z2 + x1*y2 - y1*x2 + z1*w2
    ])

def quat_conjugate(q):
    w, x, y, z = q
    return np.array([w, -x, -y, -z])

def normalize_quat(q):
    return q / np.linalg.norm(q)

def small_angle_quat(dtheta):
    angle = np.linalg.norm(dtheta)
    if angle < 1e-12:
        return np.array([1.0, 0.0, 0.0, 0.0])
    axis = dtheta / angle
    w = np.cos(angle/2.0)
    xyz = axis * np.sin(angle/2.0)
    return np.concatenate([[w], xyz])

def integrate_quat(q, omega, dt):
    dtheta = omega * dt
    dq = small_angle_quat(dtheta)
    q_new = quat_mul(dq, q)
    return normalize_quat(q_new)

def quat_to_rotmat(q):
    w, x, y, z = q
    R = np.array([
        [1-2*(y*y+z*z),   2*(x*y - z*w),  2*(x*z + y*w)],
        [2*(x*y + z*w),   1-2*(x*x+z*z),  2*(y*z - x*w)],
        [2*(x*z - y*w),   2*(y*z + x*w),  1-2*(x*x+y*y)]
    ])
    return R

def quat_to_euler(q):
    w, x, y, z = q
    sinr_cosp = 2.0*(w*x + y*z)
    cosr_cosp = 1.0 - 2.0*(x*x + y*y)
    roll = np.arctan2(sinr_cosp, cosr_cosp)
    sinp = 2.0*(w*y - z*x)
    if abs(sinp) >= 1:
        pitch = np.sign(sinp) * (np.pi/2)
    else:
        pitch = np.arcsin(sinp)
    siny_cosp = 2.0*(w*z + x*y)
    cosy_cosp = 1.0 - 2.0*(y*y + z*z)
    yaw = np.arctan2(siny_cosp, cosy_cosp)
    return np.array([roll, pitch, yaw])

def quat_to_rotvec(q):
    # returns rotation vector (3,) representing rotation from identity to q (body<-inertial)
    qn = q / np.linalg.norm(q)
    w = np.clip(qn[0], -1.0, 1.0)
    angle = 2.0 * np.arccos(w)
    if angle < 1e-12:
        return np.zeros(3)
    s = np.sqrt(1 - w*w)
    axis = qn[1:]/s
    return axis * angle

# ---------- simulation parameters ----------
dt = 0.01
T = 20.0
N = int(T / dt)
t = np.arange(N) * dt

omega_true = np.zeros((N, 3))
omega_true[:, 2] = 0.3 * np.ones(N)
omega_true[:, 0] = 0.12 * np.sin(0.6 * t)   # slightly changed motion
omega_true[:, 1] = 0.10 * np.cos(0.4 * t)

g_inertial = np.array([0.0, 0.0, 1.0])
m_inertial = np.array([0.15, 0.0, 0.9887])
m_inertial = m_inertial / np.linalg.norm(m_inertial)

gyro_bias_true = np.array([0.02, -0.015, 0.01])

gyro_noise_std = 0.004
gyro_bias_walk_std = 5e-5
accel_noise_std = 0.02
mag_noise_std = 0.01

# simulate true quaternions
qs_true = np.zeros((N,4))
q_true = np.array([1.0,0.0,0.0,0.0])
qs_true[0]=q_true
for k in range(1,N):
    q_true = integrate_quat(q_true, omega_true[k-1], dt)
    qs_true[k]=q_true

# generate measurements (with occasional magnetic disturbance and linear acceleration spike)
gyro_meas = np.zeros((N,3))
accel_meas = np.zeros((N,3))
mag_meas = np.zeros((N,3))

for k in range(N):
    qk = qs_true[k]
    Rb_i = quat_to_rotmat(qk)
    true_accel_b = Rb_i.T @ g_inertial
    true_mag_b = Rb_i.T @ m_inertial
    gyro_meas[k] = omega_true[k] + gyro_bias_true + np.random.randn(3)*gyro_noise_std

    # introduce a few linear acceleration events to test downweighting
    if 5.0 < t[k] < 6.0 or 14.5 < t[k] < 15.5:
        lin_acc = np.array([0.6*np.sin(10*t[k]), 0.4*np.cos(8*t[k]), 0.0])
    else:
        lin_acc = np.zeros(3)

    accel_vec = true_accel_b + lin_acc + np.random.randn(3)*accel_noise_std
    accel_meas[k] = accel_vec / np.linalg.norm(accel_vec)

    # introduce magnetic disturbance bursts
    if 9.8 < t[k] < 10.4 or 15.2 < t[k] < 15.4:
        mag_disturb = np.array([0.8, 0.8, -0.6])
    else:
        mag_disturb = np.zeros(3)
    mag_vec = true_mag_b + mag_disturb + np.random.randn(3)*mag_noise_std
    mag_meas[k] = mag_vec / np.linalg.norm(mag_vec)

# ---------- Robust MEKF initialization ----------
q_est = np.array([1.0,0.0,0.0,0.0])
b_est = np.zeros(3)

P = np.zeros((6,6))
P[0:3,0:3] = np.eye(3) * (np.deg2rad(3.0)**2)
P[3:6,3:6] = np.eye(3) * (0.03**2)

gyro_noise_var = gyro_noise_std**2
bias_rw_var = gyro_bias_walk_std**2
Q = np.zeros((6,6))
Q[0:3,0:3] = np.eye(3) * gyro_noise_var * dt
Q[3:6,3:6] = np.eye(3) * bias_rw_var * dt

R_acc_base = np.eye(3) * (accel_noise_std**2)
R_mag_base = np.eye(3) * (mag_noise_std**2)

# adaptive detection params
accel_norm_window = 0.5  # seconds for moving average
window_len = max(1, int(accel_norm_window/dt))
accel_norm_ma = 1.0
ma_alpha = 1.0/window_len

mag_angle_threshold = np.deg2rad(30.0)  # if measured mag deviates >30deg from predicted -> suspect
accel_norm_threshold = 0.25  # if |norm-1|>0.25 -> strong linear accel present, downweight accel
mahal_threshold = 20.0  # Mahalanobis gating threshold for 6-dim measurement

# storage
q_est_store = np.zeros((N,4))
b_est_store = np.zeros((N,3))
rotvec_errors = np.zeros((N,3))
euler_true = np.array([quat_to_euler(q) for q in qs_true])
euler_est = np.zeros((N,3))
accel_skip = np.zeros(N, dtype=bool)
mag_skip = np.zeros(N, dtype=bool)

for k in range(N):
    # PREDICTION
    omega_m = gyro_meas[k]
    omega_unbiased = omega_m - b_est
    q_est = integrate_quat(q_est, omega_unbiased, dt)

    F = np.zeros((6,6))
    F[0:3,0:3] = -skew(omega_unbiased)
    F[0:3,3:6] = -np.eye(3)
    Phi = np.eye(6) + F*dt
    P = Phi @ P @ Phi.T + Q

    # prepare measurements
    z_acc = accel_meas[k]
    z_mag = mag_meas[k]

    # predicted
    Rb_i_est = quat_to_rotmat(q_est)
    h_acc = Rb_i_est.T @ g_inertial
    h_mag = Rb_i_est.T @ m_inertial

    # detect accel linear acceleration by checking accel magnitude moving average
    accel_norm = np.linalg.norm(accel_meas[k])
    accel_norm_ma = (1-ma_alpha)*accel_norm_ma + ma_alpha*accel_norm
    accel_deviation = abs(accel_norm - 1.0)

    # set accel weight: if strong linear accel, increase R_acc (downweight)
    if accel_deviation > accel_norm_threshold:
        accel_weight = 100.0  # downweight heavily
        accel_skip[k] = True
    elif accel_deviation > (accel_norm_threshold*0.6):
        accel_weight = 10.0
    else:
        accel_weight = 1.0

    # magnetometer anomaly detection by angular difference between measured and predicted
    cosang = np.clip(np.dot(z_mag, h_mag)/(np.linalg.norm(z_mag)*np.linalg.norm(h_mag)+1e-12), -1.0, 1.0)
    mag_angle = np.arccos(cosang)
    if mag_angle > mag_angle_threshold:
        mag_weight = 100.0
        mag_skip[k] = True
    else:
        mag_weight = 1.0

    # Build H and stacked measurement
    H_acc = np.hstack([ -skew(h_acc), np.zeros((3,3)) ])
    H_mag = np.hstack([ -skew(h_mag), np.zeros((3,3)) ])
    H = np.vstack([H_acc, H_mag])
    z = np.hstack([z_acc, z_mag])
    h = np.hstack([h_acc, h_mag])

    # adaptive R based on weights
    R_block = np.zeros((6,6))
    R_block[0:3,0:3] = R_acc_base * accel_weight
    R_block[3:6,3:6] = R_mag_base * mag_weight

    # Innovation and Mahalanobis gating (robust outlier rejection)
    y = z - h
    S = H @ P @ H.T + R_block
    # prevent singular
    try:
        S_inv = np.linalg.inv(S)
    except np.linalg.LinAlgError:
        S_inv = np.linalg.pinv(S)
    maha = float(y.T @ S_inv @ y)
    # If Mahalanobis too big, attempt to drop the offending sensor from update
    if maha > mahal_threshold:
        # if mag caused large angular deviation -> remove mag (set huge R) and recompute
        if mag_angle > mag_angle_threshold:
            R_block[3:6,3:6] *= 1e3
            S = H @ P @ H.T + R_block
            try:
                S_inv = np.linalg.inv(S)
            except np.linalg.LinAlgError:
                S_inv = np.linalg.pinv(S)
            maha = float(y.T @ S_inv @ y)
        # if still too large, downweight accel as fallback
        if maha > mahal_threshold:
            R_block[0:3,0:3] *= 1e2
            S = H @ P @ H.T + R_block
            try:
                S_inv = np.linalg.inv(S)
            except np.linalg.LinAlgError:
                S_inv = np.linalg.pinv(S)
            maha = float(y.T @ S_inv @ y)
    # Final gating: if still outlier, skip update this step
    if maha > (mahal_threshold*2.0):
        # skip update entirely (no measurement correction)
        # keep covariance growth from process noise
        q_rel = quat_mul(quat_conjugate(q_est), qs_true[k])  # for error logging
        rotvec_errors[k] = quat_to_rotvec(q_rel)
        q_est_store[k] = q_est
        b_est_store[k] = b_est
        euler_est[k] = quat_to_euler(q_est)
        continue

    # compute Kalman gain and apply update
    K = P @ H.T @ np.linalg.inv(S + 1e-12*np.eye(S.shape[0]))
    dx = K @ y
    dtheta = dx[0:3]
    db = dx[3:6]
    dq = small_angle_quat(dtheta)
    q_est = normalize_quat(quat_mul(dq, q_est))
    b_est = b_est + db

    # covariance update (Joseph form)
    I6 = np.eye(6)
    P = (I6 - K @ H) @ P @ (I6 - K @ H).T + K @ R_block @ K.T

    # store
    q_est_store[k] = q_est
    b_est_store[k] = b_est
    euler_est[k] = quat_to_euler(q_est)

    # compute rotation-vector error (true relative to estimate): q_rel = q_est^{-1} * q_true
    q_est_conj = quat_conjugate(q_est)
    q_rel = quat_mul(q_est_conj, qs_true[k])
    rotvec_errors[k] = quat_to_rotvec(q_rel)

# ---------- plotting ----------
# unwrap yaw for plotting clarity
euler_true_un = np.unwrap(euler_true, axis=0)
euler_est_un = np.unwrap(euler_est, axis=0)

plt.figure(figsize=(10,5))
plt.plot(t, euler_true_un[:,0], label='True Roll')
plt.plot(t, euler_est_un[:,0], '--', label='Est Roll')
plt.plot(t, euler_true_un[:,1], label='True Pitch')
plt.plot(t, euler_est_un[:,1], '--', label='Est Pitch')
plt.plot(t, euler_true_un[:,2], label='True Yaw')
plt.plot(t, euler_est_un[:,2], '--', label='Est Yaw')
plt.xlabel('Time (s)'); plt.ylabel('Angle (rad)'); plt.title('True vs Estimated Euler Angles (unwrapped)')
plt.legend(fontsize='small'); plt.grid(True); plt.tight_layout()
plt.show()

plt.figure(figsize=(10,4))
plt.plot(t, b_est_store[:,0], '--', label='Est bias x')
plt.plot(t, b_est_store[:,1], '--', label='Est bias y')
plt.plot(t, b_est_store[:,2], '--', label='Est bias z')
plt.hlines(gyro_bias_true[0], t[0], t[-1], label='True bias x', colors='C0')
plt.hlines(gyro_bias_true[1], t[0], t[-1], label='True bias y', colors='C1')
plt.hlines(gyro_bias_true[2], t[0], t[-1], label='True bias z', colors='C2')
plt.xlabel('Time (s)'); plt.ylabel('Gyro bias (rad/s)'); plt.title('Gyro bias estimate')
plt.legend(fontsize='small'); plt.grid(True); plt.tight_layout()
plt.show()

# rotation-vector error components (rad)
plt.figure(figsize=(10,4))
plt.plot(t, rotvec_errors[:,0], label='rotvec_x (rad)')
plt.plot(t, rotvec_errors[:,1], label='rotvec_y (rad)')
plt.plot(t, rotvec_errors[:,2], label='rotvec_z (rad)')
plt.xlabel('Time (s)'); plt.ylabel('Rotation-vector error (rad)'); plt.title('Rotation-vector estimation error (body-frame)')
plt.legend(fontsize='small'); plt.grid(True); plt.tight_layout()
plt.show()

# show which steps skipped mag/accel
plt.figure(figsize=(10,2))
plt.plot(t, accel_skip.astype(int), label='accel skipped (1=yes)')
plt.plot(t, mag_skip.astype(int), label='mag skipped (1=yes)')
plt.xlabel('Time (s)'); plt.yticks([0,1]); plt.legend(fontsize='small'); plt.title('Sensor downweighting/skipping events'); plt.grid(True); plt.tight_layout()
plt.show()

# RMS on rotation-vector magnitude and individual axes
rms_rotvec = np.sqrt(np.mean(rotvec_errors**2, axis=0))
rms_rotvec_norm = np.sqrt(np.mean(np.sum(rotvec_errors**2,axis=1)))
print("RMS rotvec components (rad): x={:.4f}, y={:.4f}, z={:.4f}".format(*rms_rotvec))
print("RMS rotvec norm (rad) = {:.4f} (deg = {:.3f}°)".format(rms_rotvec_norm, rms_rotvec_norm*180/np.pi))

# summary of skipped fractions
print("Fraction accel skipped/downweighted:", np.mean(accel_skip.astype(float)))
print("Fraction mag skipped/downweighted:", np.mean(mag_skip.astype(float)))
